#!/usr/bin/env bash

# Slurm: Add nice level to jobs, or list jobs with non-zero nice level.
# Requires also the helper script "joblist".
# Homepage: https://github.com/OleHolmNielsen/Slurm_tools/

# Usage:
# Listing: $0 -l | -p partition"
# Setting: $0 [-n nice-value] [-f jobfile] job [[,]jobs ...]"
# Jobs should be comma or space separated.

# Requires also the helper script "joblist".
JOBLISTFORMAT=`which joblist`
if test $? -ne 0
then
	echo ERROR: The joblist script cannot be found
	exit 1
fi
# List or set nice values:
listnice=0
# Default nice value
NICE=-50000
SQUEUE_FORMAT=Nice,Priority,State,JobID,Partition,UserName,GroupName,StartTime,EndTime,Name,tres-per-node,Reason

# Parse command options
while getopts "lp:n:f:h" options; do
	case $options in
		l )	listnice=1
			;;
		p )	export partition="-p $OPTARG"
			echo "Select partition $OPTARG"
			listnice=1
			;;
		n )	export NICE=$OPTARG
			listnice=0
			echo Set job nice value to $NICE
			;;
		f )	export JOBFILE=$OPTARG
			if test ! -s "$JOBFILE"
			then
				echo File $JOBFILE is not a regular file or is empty
				exit 1
			fi
			listnice=0
			echo Append jobid list from file $JOBFILE
			;;
		h | * )	echo "Usage:"
			echo "Listing: $0 -l | -p partition"
			echo "Setting: $0 [-n nice-value] [-f jobfile] job [[,]jobs ...]"
			exit 1;;
	esac
done
shift $((OPTIND-1))

if test $listnice -eq 0
then
	if test -z "$JOBFILE"
	then
		JOBLIST="`$JOBLISTFORMAT $*`"
	else
		JOBLIST="`$JOBLISTFORMAT -f $JOBFILE $*`"
	fi
	if test $? -ne 0
	then
		# The $JOBLISTFORMAT script returned an error on this job list:
		echo $JOBLIST
		exit 1
	fi
	# Change nice level of jobs and print a status
	scontrol update jobid=$JOBLIST nice=$NICE
	squeue -O "$SQUEUE_FORMAT" -j $JOBLIST
else
	echo; echo "Jobs with a non-zero nice value:"; echo
	squeue $partition -O "$SQUEUE_FORMAT" | awk '{if ($1 != 0) print $0}' 
fi
